#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Sun Dec 29 14:32:57 2019

@author: hyungmokson
"""

import csv
import numpy as np
from os.path import isfile, join

## Header index:
## 0: Measurement time
## 1: Variable
## 2: STIRAP freq 1 (THz)
## 3: STIRAP freq 2 (THz)
## 4: STIRAP freq 2 (THz)

filename = 'double evap_with Na.csv'
with open(filename, 'r') as csvfile:
    csvreader = csv.reader(csvfile)
    header = next(csvreader)

    ncols = len(csvfile.readline().split(','))

#print("header = " + str(header))

list_cols = list(range(1, ncols))
list_cols.remove(2)
list_cols.remove(3)
list_cols.remove(4) ## remove STIRAP frequency entries
raw = np.loadtxt(filename, delimiter = ',', skiprows=1, usecols = tuple(list_cols))

## Import average fit err factor (delta_T/T) from seperate txt file
file_fiterr_frac_avg = 'mol temp_fit err frac.txt'
fiterr_frac_avg = np.loadtxt(file_fiterr_frac_avg, delimiter = '\t')
fiterr_avg = {}
for e in fiterr_frac_avg:
    key = e[0]
    if fiterr_avg.get(key, -1) == -1:
        fiterr_avg[key] = e[1]
    else:
        fiterr_avg[key].append(e[1])

## create hashmap for the raw data 
data_idx_start = 1
data = {}
for row in raw:
    key = row[0]
    if data.get(key, -1) == -1:
        data[key] = [row[data_idx_start:]]
    else:
        data[key].append(row[data_idx_start:])

keys = list(data.keys())
arg_list = np.argsort(keys)

data_array = [] 
i = 0
## first column na num
## 4th column, the last one is psd factor
for key, value in data.items():
    n = len(data[key])
    print("for " + str(key) + ":  " + str(n))
    temp = np.array(value)

    avg_array = np.average(temp, axis = 0)[1:3]
    std_array = np.std(temp, axis = 0)[1:3]

    ## last one is PSD factor, ignore it
    stand_err_array = std_array/np.sqrt(n-1) 

    ## only statistical error is considered.    
    avg_data_array = avg_array
    err_array = stand_err_array

    ## fit error is calculated as the standard error of the averages
#    stand_err_rms = fiterr_avg[key] * avg_array[1]/np.sqrt(n)
    stand_err_rms = fiterr_avg[key] * avg_array[1]
    moltemp_err = np.sqrt(stand_err_array[1]**2 + stand_err_rms**2)

    err_array = [stand_err_array[0], moltemp_err]
    ## array for each "Variable" with avgerage data and total error
#    final_array = []
#    q = np.insert(avg_data_array, 0, key)
    final_array = np.append(avg_data_array, err_array)

    ## append into the final datay array
    data_array.append(final_array)    
    i += 1


output = []
## sort by the variables in ascending order
for idx in arg_list:
    temp = np.insert(data_array[idx], 0, keys[idx])
    output.append(temp)
    
idx_molnum = 1
with open(join('csv','mol num.csv'), 'w') as file:
    writer = csv.writer(file, delimiter = ',')
    for row in output:
        writer.writerow([row[0], row[idx_molnum], row[idx_molnum+2]])

idx_moltemp = 2
with open(join('csv', 'mol temp.csv'), 'w') as file:
    writer = csv.writer(file, delimiter = ',')
    for row in output:
        writer.writerow([row[0], row[idx_moltemp], row[idx_moltemp+2]])


        